#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import PIL
import glob
import pycdlib
import pathlib
import subprocess
import pytesseract
from .._util import Util, FloppyMount, VirtualUser
from .._errors import BuildError
from .._workdir import WorkDirAddonBase
from .._vm import Vm
from ._spec import Spec
from ._batch import BatchFile


class WorkDirAddon(WorkDirAddonBase):

    def __init__(self, parent):
        self._parent = parent

    @property
    def installation_iso_filepath(self):
        return os.path.join(self._parent._path, "install.iso")

    @property
    def extracted_floppy_filepath(self):
        return os.path.join(self._parent._path, "extracted.ima")

    @property
    def boot_floppy_filepath(self):
        return os.path.join(self._parent._path, "boot.ima")

    @property
    def bios_filepath(self):
        return os.path.join(self._parent._path, "bios.bin")

    @property
    def patcher9x_floppy_filepath(self):
        return os.path.join(self._parent._path, "patcher9x.ima")

    @property
    def simd95_filepath(self):
        return os.path.join(self._parent._path, "simd95.com")

    @property
    def drivers_dirpath(self):
        return os.path.join(self._parent._path, "drivers")

    @property
    def floppy_op_dirpath(self):
        return os.path.join(self._parent._path, "floppy.dir")

    @property
    def iso_op_dirpath(self):
        return os.path.join(self._parent._path, "iso.dir")

    @property
    def tmp_file(self):
        return os.path.join(self._parent._path, "tmp.file")


class BuildUtil:

    @staticmethod
    def prepare(ts, workDirObj, srcMedia):
        # extract needed iso files
        subprocess.check_call(["7z", "x", "-o" + workDirObj.iso_op_dirpath, srcMedia.get_iso_filepath(), "setup.exe", "win98", "drivers"], stdout=subprocess.DEVNULL)
        for fn in os.listdir(os.path.join(workDirObj.iso_op_dirpath, "win98")):
            if not fn.endswith(".CAB") and fn.lower() != fn:
                os.rename(os.path.join(workDirObj.iso_op_dirpath, "win98", fn), os.path.join(workDirObj.iso_op_dirpath, "win98", fn.lower()))        # not neccessary, only to make file names prettier

        # create real install iso
        if True:
            # add batch file
            batch_file = BatchFile(ts)
            with open(os.path.join(workDirObj.iso_op_dirpath, batch_file.get_filepath()), "wb") as f:
                f.write(batch_file.get_binary_content())

            # extract and modify boot floppy image
            subprocess.check_call(["mformat", "-i", workDirObj.boot_floppy_filepath, "-f", "1440", "-C"])
            if srcMedia.get_boot_floppy_filepath() is not None:
                with FloppyMount(srcMedia.get_boot_floppy_filepath(), workDirObj.floppy_op_dirpath):
                    for fullfn in glob.glob(os.path.join(workDirObj.floppy_op_dirpath, "*")):
                        subprocess.check_call(["mcopy", "-i", workDirObj.boot_floppy_filepath, fullfn, "::"])
                Util.floppyCopyBootstrapCode(srcMedia.get_boot_floppy_filepath(), workDirObj.boot_floppy_filepath)
            else:
                with open(workDirObj.extracted_floppy_filepath, "wb") as f:
                    f.write(Util.isoExtractEltoritoFloppyImageData(srcMedia.get_iso_filepath()))
                with FloppyMount(workDirObj.extracted_floppy_filepath, workDirObj.floppy_op_dirpath):
                    # copy files
                    for fullfn in glob.glob(os.path.join(workDirObj.floppy_op_dirpath, "*")):
                        if os.path.basename(fullfn) == "CONFIG.SYS":
                            with open(workDirObj.tmp_file, "wb") as f:
                                buf = pathlib.Path(fullfn).read_bytes()
                                buf = buf.replace(b'menudefault=SETUP_CD,30', b'menudefault=SETUP_CD,1')    # we always use the default option, make it faster
                                f.write(buf)
                            subprocess.check_call(["mcopy", "-i", workDirObj.boot_floppy_filepath, workDirObj.tmp_file, "::CONFIG.SYS"])
                        else:
                            subprocess.check_call(["mcopy", "-i", workDirObj.boot_floppy_filepath, fullfn, "::"])
                Util.floppyCopyBootstrapCode(workDirObj.extracted_floppy_filepath, workDirObj.boot_floppy_filepath)

            # make iso
            iso = pycdlib.PyCdlib()
            iso.new(joliet=3)
            try:
                # inject files
                for root, dirs, files in os.walk(workDirObj.iso_op_dirpath):
                    root2 = root.replace(workDirObj.iso_op_dirpath, "")
                    if root2 == "":
                        root2 = "/"
                    for dn in dirs:
                        Util.pycdlibAddDirectory(iso, joliet_path=os.path.join(root2, dn))
                    for fn in files:
                        Util.pycdlibAddFile(iso, pathlib.Path(os.path.join(root, fn)).read_bytes(), joliet_path=os.path.join(root2, fn))

                # inject boot image
                ret = Util.pycdlibAddFile(iso, pathlib.Path(workDirObj.boot_floppy_filepath).read_bytes(), joliet_path="/boot.ima")
                iso.add_eltorito(bootfile_path=ret["iso_path"], media_name="floppy")

                iso.write(workDirObj.installation_iso_filepath)
            finally:
                iso.close()

        # get no pnp bios bin file
        subprocess.check_call(["cp", srcMedia.get_no_pnp_bios_filepath(), workDirObj.bios_filepath])

        # get patcher9x floppy image
        subprocess.check_call(["cp", srcMedia.get_patcher9x_floppy_filepath(), workDirObj.patcher9x_floppy_filepath])

        # get simd95.com
        if ts.simd95:
            subprocess.check_call(["cp", srcMedia.get_simd95_filepath(), workDirObj.simd95_filepath])

        # drivers from installation iso file
        os.makedirs(workDirObj.drivers_dirpath)
        if True:
            driverSet = set()

            # extract driver packs
            pass

            # extract driver packs
            for name, fullfn in srcMedia.get_driverpack_filepath_dict().items():
                if name not in ts.extra_drivers:
                    continue
                curDir = os.path.join(workDirObj.drivers_dirpath, name)
                os.makedirs(curDir)
                if fullfn.endswith(".zip") or fullfn.endswith(".7z") or fullfn.endswith(".iso"):
                    subprocess.check_call(["7z", "x", "-o" + curDir, fullfn], stdout=subprocess.DEVNULL)
                    while Util.moveSoleDirContentOut(curDir):
                        pass
                else:
                    assert False
                driverSet.add(name)

            assert driverSet == set(ts.extra_drivers)

    @staticmethod
    def osInstall(ts, workDirObj, bDebug):
        spec = Spec()

        assert workDirObj.get_image_file_storage_layout_name() == spec.getStorageLayoutName()

        # install-1: collect information, copy files
        vm = Vm(workDirObj)
        vm.set_record_file("install-1.mp4")
        vm.set_no_reboot(True)
        vm.use_cdrom(workDirObj.installation_iso_filepath, boot=True)
        vm.start(show=bDebug)
        try:
            vuser = VirtualUser98("eng", vm)
            vuser.operate(vuser.BootFromCdrom())
            vuser.sleep(30)
            vuser.operate(vuser.InitConfirm())
            vuser.operate(vuser.InitConfirm2())
            vm.wait_until_stop()
        except BaseException:
            vm.stop()
            raise

        # patcher9x
        vm = Vm(workDirObj)
        vm.set_record_file("patcher9x.mp4")
        vm.use_floppy(0, workDirObj.patcher9x_floppy_filepath, boot=True)
        vm.start(show=bDebug)
        try:
            vuser = VirtualUser98("eng", vm)
            vuser.sleep(10)
            vuser.operate(vuser.Patcher9xRun())
            vuser.operate(vuser.Patcher9xShutdown())
            vm.wait_until_stop()
        except BaseException:
            vm.stop()
            raise

        # install-2: confirm EULA, detect hardware
        vm = Vm(workDirObj)
        vm.set_record_file("install-2.mp4")
        vm.set_no_reboot(True)
        vm.use_custom_bios(workDirObj.bios_filepath)                    # win98 wrongly recognizes the "PCI bus" as "PnP BIOS", so we use a no-pnp-bios until PCI bus is detected and recorded
        vm.start(show=bDebug)
        try:
            # we don't use msbatch.inf to skip the following 4 steps, because we don't want to calculate ProductId (xxxxx-OEM-xxxxxxx-xxxxx)
            vuser = VirtualUser98("eng", vm)
            vuser.sleep(10)
            vuser.operate(vuser.FirstBootUserAndOrg())
            vuser.operate(vuser.FirstBootEula())
            vuser.operate(vuser.FirstBootProductKey())
            vuser.operate(vuser.FirstBootFinish())
            vm.wait_until_stop()
        except BaseException:
            vm.stop()
            raise

        # process presave and postsave
        if True:
            pass

        # install-3: setup Control Panel, Program on the Start menu, etc.
        vm = Vm(workDirObj)
        vm.set_record_file("install-3.mp4")
        vm.set_no_reboot(True)
        vm.use_custom_bios(workDirObj.bios_filepath)                    # win98 wrongly recognizes the "PCI bus" as "PnP BIOS", so we use a no-pnp-bios until PCI bus is detected and recorded
        vm.start(show=bDebug)
        try:
            vm.wait_until_stop()
        except BaseException:
            vm.stop()
            raise

        # copy simd95.com
        if ts.simd95:
            with workDirObj.mount_image_file(stage=1) as mp:
                subprocess.check_call(["mcopy", "-i", mp.drive_c_loop_file, workDirObj.simd95_filepath, "::simd95.com"])
                # FIXME: modify autoexec.bat

        # copy extra drivers
        with workDirObj.mount_image_file(stage=1) as mp:
            # subprocess.check_call(["mcopy", "-s", "-i", mp, workDirObj.drivers_dirpath, "::DRIVERS"])
            fset = set()
            # FIXME: put files already in PRECOPY dir into fset
            for fn in os.listdir(workDirObj.drivers_dirpath):
                fullfn = os.path.join(workDirObj.drivers_dirpath, fn)                                                         # every driver directory
                for fn2 in os.listdir(fullfn):
                    fullfn2 = os.path.join(fullfn, fn2)                                                                       # every file and subdirectory in driver directory
                    assert fullfn2 not in fset                                                                                # win98 doesn't search driver directory recursively, not only *.inf but also other files that *.inf refers, we must make all files flatten and check no file has same name
                    subprocess.check_call(["mcopy", "-s", "-i", mp.drive_c_loop_file, fullfn2, "::WINDOWS/SYSTEM/PRECOPY"])   # win98 searches C:\WINDOWS\SYSTEM\PRECOPY for driver files by default no matter what value DevicePath and OtherDevicePath have in registry, we have to copy files to this directory to comply with this behavior
                    fset.add(fullfn2)

        # install-final: use standard hardware, install drivers
        if True:
            vm = Vm(workDirObj)
            vm.set_record_file("install-final.mp4")
            vm.start(show=bDebug)
            try:
                vuser = VirtualUser98("eng", vm)
                vuser.sleep(5)
                vuser.operate(vuser.FinalBootDriverWizard())
                vuser.operate(vuser.FinalBootWaitStartMenu())
                vm.wait_until_stable()
                vuser.operate(vuser.FinalBootShutdown())
                vm.wait_until_stop()
            except BaseException:
                vm.stop()
                raise

    @staticmethod
    def cleanup(ts, workDirObj):
        pass


class VirtualUser98(VirtualUser):

    class BootFromCdrom:

        def do(self, parent, img):
            # OCR result should be:
            #
            # Microsoft Windows 98 CD-ROM Startup Menu
            #
            # Boot from Hard Disk
            #
            # Boot from CD-ROM
            #
            # Enter your choice: _ Time remaining: 10
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Boot from CD-ROM" in s and "Enter your choice:" in s:
                parent._vm.input_keyboard("keys", "2")
                return True
            return False

    class InitConfirm:

        def do(self, parent, img):
            # OCR result should be:
            #
            # Microsoft Windows 98 Setup
            #
            # Welcome to Setup.
            #
            # The Setup program prepares Windows 98 to run on your
            # computer.
            #
            # + To set up Windows now, press ENTER.
            # + To learn more about Setup before continuing, press Fi.
            # + To quit Setup without installing Windows, press F3.
            # Note: If you have not backed up your files recently, you
            # might want to do so before installing Windows. To back
            # up your files, press F3 to quit Setup now. Then, back
            #
            # up your files by using a backup program
            #
            # To continue with Setup, press ENTER.
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Welcome to Setup." in s and "To continue with Setup, press ENTER." in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class InitConfirm2:

        def do(self, parent, img):
            # OCR result should be:
            #
            # Setup is preparing to install Windows.
            # Please wait while Setup initializes.
            # Setup is now going to perform a routine check on your system
            #
            # To continue, press ENTER. To quit Setup, press ESC.
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Setup is preparing to install Windows." in s and "To continue, press ENTER." in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class Patcher9xRun:

        def do(self, parent, img):
            # OCR result should be:
            #
            # HSUCDX by Jason Hood < jadoxa@yahoo.com.au>.
            # Version 3.08 (1? February, 2021). Freeware.
            #
            # Derived from v1.4b by
            # John H. McCoy, October 2000,
            #
            # http://shsucdx.adoxa .vze.com/ 1 Sam Houston State University.
            #
            # Can’t open CD driver SHSU-CDX. SHSUCDX can’t install.
            # FreeCom version @.85a - WATCOMC - XMS_Suap [Jul 10 2021 19:28:061
            #
            # -- BOOT SUCCESS
            #
            # Bootable floppy with patch for Windous 95/98/Me
            # for run on newest CPUs - AMD ZEN 2+, Intel Tiger Lake+
            #
            # Yersion: 0.8.50-BETA4
            #
            # For start patching (interactive mode) type:
            # patch9x
            #
            # For view help type:
            # patch9x -h
            #
            # For info about other tools on this disk type:
            # info
            # >
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "BOOT SUCCESS" in s:
                parent._vm.input_keyboard("line", r"patch9x.exe -auto -no-backup C:\WINDOWS\SYSTEM")
                return True
            return False

    class Patcher9xShutdown:

        def do(self, parent, img):
            # OCR result should contain (extra garbage may exist):
            #
            # Status: new patches to apply: 2
            # Patch applied successfully
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Patch applied successfully" in s:
                parent._vm.poweroff()          # DOS does not support shutdown
                return True
            return False

    class FirstBootUserAndOrg:

        def do(self, parent, img):
            # OCR result should contain (extra garbage may exist):
            #
            # User Information
            #
            # Type your name below. Ifyou want, you can alsa type the name
            # ofthe company you work for.
            #
            # Name:
            #
            # Company:
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "User Information" in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class FirstBootEula:

        def do(self, parent, img):
            # OCR result should contain (extra garbage may exist):
            #
            # MICROSOFT WINDOWS 98 SECOND EDITION
            #
            # END-USER LICENSE AGREEMENT FOR
            # MICROSOFT DESKTOP OPERATING SYSTEMS
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "END-USER LICENSE AGREEMENT" in s:
                parent._vm.input_keyboard("keys", "alt+a", "kp_enter")
                return True
            return False

    class FirstBootProductKey:

        def do(self, parent, img):
            # OCR result should contain (extra garbage may exist):
            #
            # Windows Product Key
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Windows Product Key" in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class FirstBootFinish:

        def do(self, parent, img):
            # OCR result should contain (extra garbage may exist):
            #
            # Start Wizard
            #
            # Windows 98 saved all information,
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Start Wizard" in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class FinalBootDriverWizard:

        def __init__(self):
            selfDir = os.path.dirname(os.path.realpath(__file__))
            self._imgLeftTop = PIL.Image.open(os.path.join(selfDir, "final-left-top.png"))
            self._imgLeftBottom = PIL.Image.open(os.path.join(selfDir, "final-left-bottom.png"))

        def do(self, parent, img):
            if Util.imgMatchSubImg(img, self._imgLeftTop, "left-top") and Util.imgMatchSubImg(img, self._imgLeftBottom, "left-bottom"):
                # this operation is finished
                return True

            img = Util.imgCropPureColorBorder(img)                                                                      # pytesseract recognizes nothing if left top corner is empty, sigh
            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            print(s)

            # OCR result should contain (extra garbage may exist):
            #
            # This wizard searches for new drivers for:
            #
            # XXXXXXXXX
            #
            # A device driver is a software program that makes a
            # hardware device work.
            #
            if "This wizard searches for new " in s:                                                                    # pyteesseract recognize "drivers" as "civers", sigh
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            # OCR result should contain (extra garbage may exist):
            #
            # What do you want Windows to do?
            #
            if "What do you want Windows to do" in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            # OCR result should contain (extra garbage may exist):
            #
            # Windows will search for new drivers in its driver database
            # on your hard drive, and in any of the following selected
            # locations. Cick Next to start the search.
            #
            if re.search(r"\S+ \S+ search for new ", s, re.M) is not None and " any of the following selected" in s:    # pytessearact result has many errors, sigh
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            # OCR result should contain (extra garbage may exist):
            #
            # Windows driver file search for the device:
            #
            # XXXXXXXXX
            #
            # Windows is now ready to install the best driver for this
            # device. Click Back to select a different driver, or cick Next
            # to continue.
            #
            # Location of driver:
            #
            if re.search(r"Win\S+ \S+ \S+ search \S+ \S+ \S+:", s, re.M) is not None:                                   # pytessearact result has many errors, sigh
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            # OCR result should contain (extra garbage may exist):
            #
            # Windows has finished installing the software that your new
            # hardware device requires.
            #
            if "software that your new" in s and "hardware device requires" in s:                                       # pytessearact result has many errors, sigh
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            # OCR result should contain (extra garbage may exist):
            #
            # Do you want to restart your computer now?
            #
            if re.search("Do you want to .* your computer now", s, re.M) is not None:                                   # pytessearact result has many errors, sigh
                parent._vm.input_keyboard("keys", "kp_enter")
                return False

            return False

    class FinalBootWaitStartMenu:

        def __init__(self):
            selfDir = os.path.dirname(os.path.realpath(__file__))
            self._imgLeftTop = PIL.Image.open(os.path.join(selfDir, "final-left-top.png"))
            self._imgLeftBottom = PIL.Image.open(os.path.join(selfDir, "final-left-bottom.png"))

        def do(self, parent, img):
            if not Util.imgMatchSubImg(img, self._imgLeftTop, "left-top"):
                return False
            if not Util.imgMatchSubImg(img, self._imgLeftBottom, "left-bottom"):
                return False
            return True

    class FinalBootShutdown:

        def do(self, parent, img):
            # win98 does not respond to the QMP command "system_powerdown", so we use UI operation instead of vm.shutdown()
            parent._vm.input_keyboard("keys", "alt+f4", "kp_enter")
            return True

    def __init__(self, ocr_lang, vm):
        super().__init__(vm, BuildError)
        self._ocrLang = ocr_lang
